package rabbitmq.publisher.confirms;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import rabbitmq.constant.Constants;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;

public class PublisherConfirms {

    private static final Integer MESSAGE_COUNT = 10000;
    static Connection createConnection() throws Exception {
        ConnectionFactory connectionFactory = new ConnectionFactory();
        connectionFactory.setHost(Constants.HOST);
        connectionFactory.setPort(Constants.PORT); //需要提前开放端口号
        connectionFactory.setUsername(Constants.USER_NAME);//账号
        connectionFactory.setPassword(Constants.PASSWORD);  //密码
        connectionFactory.setVirtualHost(Constants.VIRTUAL_HOST); //虚拟主机
        return connectionFactory.newConnection();
    }

    public static void main(String[] args) throws Exception {

        // Strategy #1: Publishing Messages Individually
        // 单独确认
        publishingMessagesIndividually();

        // Strategy #2: Publishing Messages in Batches
        // 批量确认
        publishingMessagesInBatches();

        // Strategy #3: Handling Publisher Confirms Asynchronously
        // 异步确认
        handlingPublisherConfirmsAsynchronously();

    }

    /**
     * 单独确认
     */
    private static void publishingMessagesIndividually() throws Exception {

        try(Connection connection = createConnection()) {

            // 1. 开启信道
            Channel channel = connection.createChannel();

            // 2. 设置信道为confirm模式
            channel.confirmSelect();

            // 3. 声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE1, true, false, false, null);

            // 4. 发送消息, 并等待确认
            long start = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {

                String msg = "hello publisher confirms" + i;
                channel.basicPublish("",Constants.PUBLISHER_CONFIRMS_QUEUE1, null, msg.getBytes());

                // 等待 broker 确认
                // 最长等待时间 5 s
                channel.waitForConfirmsOrDie(5000);
            }
            long end = System.currentTimeMillis();
            System.out.printf("单独确认策略, 消息条数 : %d, 耗时 : %d ms \n",MESSAGE_COUNT, end - start);
        }
    }

    /**
     * 批量确认
     * @throws Exception
     */
    private static void publishingMessagesInBatches() throws Exception{

        try(Connection connection = createConnection()) {

            // 1. 开启信道
            Channel channel = connection.createChannel();

            // 2. 设置信道为confirm模式
            channel.confirmSelect();

            // 3. 声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE2, true, false, false, null);

            // 4. 发送消息, 并进行确认
            long start = System.currentTimeMillis();

            // 指定一次性处理的消息
            int batchSize = 100;
            int outstandingMessageCount = 0;

            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "hello publisher confirms" +  i;
                channel.basicPublish("",Constants.PUBLISHER_CONFIRMS_QUEUE2, null, msg.getBytes());
                outstandingMessageCount++;

                // 当消息数量到达指定一次处理的消息数的时候, 批量一起处理
                if (outstandingMessageCount == batchSize){
                    channel.waitForConfirmsOrDie(5000);
                    outstandingMessageCount = 0;
                }
            }

            // 将最后没有到指定的消息数的消息, 一起处理
            if (outstandingMessageCount > 0){
                channel.waitForConfirmsOrDie(5000);
            }
            long end = System.currentTimeMillis();
            System.out.printf("批量确认策略, 消息条数 : %d, 耗时 : %d ms \n",MESSAGE_COUNT, end - start);
        }
    }

    /**
     * 异步确认
     */
    private static void handlingPublisherConfirmsAsynchronously() throws Exception{

        try (Connection connection = createConnection()){

            // 1. 开启信道
            Channel channel = connection.createChannel();

            // 2. 设置信道为confirm模式
            channel.confirmSelect();

            // 3. 声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE3, true, false, false, null);

            // 4. 监听confirm
            long start = System.currentTimeMillis();

//            // 有序集合, 存储未确认的消息 ID
//            SortedSet<Long> confirmSeqNo = Collections.synchronizedSortedSet(new TreeSet<>());

            // 存放 deliveryTag 和 msg 的集合
            // key 为 deliveryTag，value 为对应的 msg
            /**
             * 选择 ConcurrentNavigableMap 而不是 ConcurrentHashMap：
             * 基于跳表实现，清除所有 <= deliveryTag 的时间复杂度为 O(log n)
             * 线程安全
             * 保证了原子性
             */
            ConcurrentNavigableMap<Long, String> map = new ConcurrentSkipListMap<>();

            channel.addConfirmListener(new ConfirmListener() {

                /**
                 * 消息成功确认
                 * @param deliveryTag : 表示当前消息
                 * @param multiple : true 表示批量删除；false 表示值删除当前 deliveryTag
                 */
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {

                    // confirmSeqNo.headSet(n) : 返当前集中小于 n 的集合
                    if (multiple){

                        // 批量确认 : 将集合中⼩于等于当前序号 deliveryTag 元素的集合清除,表⽰
                        // 这批序号的消息都已经被 ack 了
//                        confirmSeqNo.headSet(deliveryTag + 1).clear();

                        // true : 表示包含了 deliveryTag
                        map.headMap(deliveryTag, true).clear();
                    }else {
                        // 单条确认 : 将当前的 deliveryTag 从集合中移除
//                        confirmSeqNo.remove(deliveryTag);
                        map.remove(deliveryTag);
                    }
                }

                /**
                 * 消息确认失败
                  * @param deliveryTag
                 * @param multiple
                 */
                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {

                    if (multiple) {
                        // 批量处理：获取所有需要重发的消息
                        ConcurrentNavigableMap<Long, String> failedMessages = map.headMap(deliveryTag, true);

                        // 先保存需要重发的消息内容
                        Map<Long, String> toResend = new HashMap<>(failedMessages);

                        // 清除失败记录
                        failedMessages.clear();

                        // 重新发送所有失败的消息
                        for (String msg : toResend.values()) {
                            long newDeliveryTag = channel.getNextPublishSeqNo();
                            channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE3, null, msg.getBytes(StandardCharsets.UTF_8));
                            map.put(newDeliveryTag, msg);
                        }
                    } else {
                        // 单条处理：获取失败消息内容
                        String msg = map.get(deliveryTag);

                        // 清除失败记录
                        map.remove(deliveryTag);

                        // 重新发送
                        if (msg != null) {
                            long newDeliveryTag = channel.getNextPublishSeqNo();
                            channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE3, null, msg.getBytes(StandardCharsets.UTF_8));
                            map.put(newDeliveryTag, msg);
                        }
                    }
                }
            });

            // 5. 发送消息
            for (int i = 0; i < MESSAGE_COUNT; i++) {

                String msg = "hello publisher confirms" + i;

                // 获取该条消息的唯一序列号, 并为下一条消息分配唯一序列号
                long seqNo = channel.getNextPublishSeqNo();
                channel.basicPublish("",Constants.PUBLISHER_CONFIRMS_QUEUE3, null, msg.getBytes());

                // 将序号存入集合中
//                confirmSeqNo.add(seqNo);
                map.put(seqNo, msg);

            }

            // 保证消息都已经确认
//            while (!confirmSeqNo.isEmpty()){
//                Thread.sleep(10);
//            }
            while (!map.isEmpty()) {
                Thread.sleep(10);
            }

            long end = System.currentTimeMillis();
            System.out.printf("异步确认策略, 消息条数: %d, 耗时: %d ms \n",MESSAGE_COUNT, end - start);
        }
    }
}
